# -*- coding: utf-8 -*- #
"""*********************************************************************************************"""
#   FileName     [ asr/ctc.py ]
#   Synopsis     [ ctc loss for asr]
#   Author       [ Andy T. Liu (Andi611) ]
#   Copyright    [ Copyleft(c), Speech Lab, NTU, Taiwan ]
#   Reference 1  [ https://github.com/Alexander-H-Liu/End-to-end-ASR-Pytorch ]
"""*********************************************************************************************"""


###############
# IMPORTATION #
###############
import numpy as np


''' 
    CTC Prefix score calculator
    An implementation of Algo. 2 in https://www.merl.com/publications/docs/TR2017-190.pdf (Watanabe et. al.)
    Reference (official implementation): https://github.com/espnet/espnet/tree/master/espnet/nets
'''
class CTCPrefixScore():
    def __init__(self, x):
        self.logzero = -100000000.0
        self.blank = 0
        self.eos = 1
        self.x = x.cpu().numpy()[0]
        self.odim = x.shape[-1]
        self.input_length = len(self.x)


    def init_state(self):
        # 0 = non-blank, 1 = blank
        r = np.full((self.input_length,2),self.logzero,dtype=np.float32)
        
        # Accumalate blank at each step
        r[0,1] = self.x[0,self.blank]
        for i in range(1,self.input_length):
            r[i,1] = r[i-1,1] + self.x[i,self.blank]
        return r
    
    def full_compute(self,g,r_prev):
        '''Given prefix g, return the probability of all possible sequence y (where y = concat(g,c))
           This function computes all possible tokens for c (memory inefficient)'''
        prefix_length = len(g)
        last_char = g[-1] if prefix_length>0 else 0

        # init. r
        r = np.full((self.input_length,2,self.odim),self.logzero,dtype=np.float32)
        
        # start from len(g) because is impossible for CTC to generate |y|>|X|
        start = max(1,prefix_length)

        if prefix_length == 0:
            r[0,0,:] = self.x[0,:]    # if g = <sos>
        
        psi = r[start-1,0,:]
        
        phi = np.logaddexp(r_prev[:,0],r_prev[:,1])
        
        for t in range(start,self.input_length):
            # prev_blank
            prev_blank = np.full((self.odim),r_prev[t-1,1],dtype=np.float32)
            # prev_nonblank
            prev_nonblank = np.full((self.odim),r_prev[t-1,0],dtype=np.float32)
            prev_nonblank[last_char] = self.logzero
            
            phi = np.logaddexp(prev_nonblank, prev_blank)
            # P(h|current step is non-blank) = [ P(prev. step = y) + P()]*P(c)
            r[t,0,:] = np.logaddexp(r[t-1,0,:],phi) + self.x[t,:]
            # P(h|current step is blank) = [P(prev. step is blank) + P(prev. step is non-blank)]*P(now=blank)
            r[t,1,:] = np.logaddexp(r[t-1,1,:],r[t-1,0,:]) + self.x[t,self.blank] 
            psi = np.logaddexp(psi,phi+self.x[t,:])

        #psi[self.eos] = np.logaddexp(r_prev[-1,0], r_prev[-1,1])
        return psi,np.rollaxis(r, 2)
    
    def cheap_compute(self,g,r_prev,candidates):
        '''Given prefix g, return the probability of all possible sequence y (where y = concat(g,c))
           This function considers only those tokens in candidates for c (memory efficient)'''
        prefix_length = len(g)
        odim = len(candidates)
        last_char = g[-1] if prefix_length>0 else 0

        # init. r
        r = np.full((self.input_length,2,len(candidates)),self.logzero,dtype=np.float32)
        
        # start from len(g) because is impossible for CTC to generate |y|>|X|
        start = max(1,prefix_length)

        if prefix_length == 0:
            r[0,0,:] = self.x[0,candidates]    # if g = <sos>
        
        psi = r[start-1,0,:]
        
        phi = np.logaddexp(r_prev[:,0],r_prev[:,1])
        
        for t in range(start,self.input_length):
            # prev_blank
            prev_blank = np.full((odim),r_prev[t-1,1],dtype=np.float32)
            
            # prev_nonblank
            prev_nonblank = np.full((odim),r_prev[t-1,0],dtype=np.float32)
            if last_char in candidates:
                prev_nonblank[candidates.index(last_char)] = self.logzero

            phi = np.logaddexp(prev_nonblank, prev_blank)
            # P(h|current step is non-blank) = [ P(prev. step = y) + P()]*P(c)
            r[t,0,:] = np.logaddexp(r[t-1,0,:],phi) + self.x[t,candidates]
            # P(h|current step is blank) = [P(prev. step is blank) + P(prev. step is non-blank)]*P(now=blank)
            r[t,1,:] = np.logaddexp(r[t-1,1,:],r[t-1,0,:]) + self.x[t,self.blank] 
            psi = np.logaddexp(psi,phi+self.x[t,candidates])

        psi[self.eos] = np.logaddexp(r_prev[-1,0], r_prev[-1,1]) 
        return psi,np.rollaxis(r, 2)

    
